local bit = require "bit"
local cjson = require "cjson"
local sunpack = string.unpack or require "compat53.string".unpack

local awdb_functions = {}
local awdb_meta = {
    __name = "awdb-lua";
    __index = awdb_functions;
}

-- Data type parsing functions
local data_parse_type_functions = {}
local parse_type_func = {}

-- parse array
function awdb_functions:parse_array(offset)
    local data_len = tonumber(self.conts:byte(offset + 1))
    offset = offset + 1
    local data = {}
    local value = nil
    for i = 1, data_len do
        value, offset = self:parse_content(self, offset)
        data[i] = value
    end
    return data, offset
end

-- parse pointer
function awdb_functions:parse_pointer(offset)
    local data_len = tonumber(self.conts:byte(offset + 1))
    offset = offset + 1
    local temp_offset = offset + data_len
    local str = self.conts:sub(offset + 1, temp_offset)
    local buf = sunpack(">I" .. data_len, str)
    local base_offset = tonumber(self.meta["base_offset"])
    local pointer = base_offset + buf
    local value, _ = self:parse_content(self, pointer)
    return value, temp_offset
end

-- parse string
function awdb_functions:parse_string(offset)
    local data_len = tonumber(self.conts:byte(offset + 1))
    offset = offset + 1
    local temp_offset = offset + data_len
    return self.conts:sub(offset + 1, temp_offset), temp_offset
end

-- parse long string
function awdb_functions:parse_long_string(offset)
    local size = tonumber(self.conts:byte(offset + 1))
    offset = offset + 1
    local str = string.sub(self.conts, offset + 1, offset + size)
    local data_len = sunpack(">i" .. size, str)
    offset = offset + size
    local temp_offset = offset + data_len
    return self.conts:sub(offset + 1, temp_offset), temp_offset
end

-- function to parse unsigned int
function awdb_functions:parse_unsigned(offset)
    local data_len = tonumber(self.conts:byte(offset + 1))
    offset = offset + 1
    local temp_offset = offset + data_len
    local str = string.sub(self.conts, offset + 1, temp_offset)
    local value = sunpack(">I" .. data_len, str)
    return value, temp_offset
end

-- function to parse signed int
function awdb_functions:parse_signed(offset)
    local temp_offset = offset + 4
    local str = string.sub(self.conts, offset + 1, temp_offset)
    local value = sunpack(">i", str)
    return value, temp_offset
end

-- function to parse float
function awdb_functions:parse_float(offset)
    local temp_offset = offset + 4
    local str = string.sub(self.conts, offset + 1, temp_offset)
    local value = sunpack(">f", str)
    return value, temp_offset
end

-- function to parse double
function awdb_functions:parse_double(offset)
    local temp_offset = offset + 8
    local str = string.sub(self.conts, offset + 1, temp_offset)
    local value = sunpack(">d", str)
    return value, temp_offset
end

-- parse array
data_parse_type_functions[1] = awdb_functions.parse_array
-- parse pointer
data_parse_type_functions[2] = awdb_functions.parse_pointer
-- parse string
data_parse_type_functions[3] = awdb_functions.parse_string
-- parse long string
data_parse_type_functions[4] = awdb_functions.parse_long_string
-- parse unsigned int
data_parse_type_functions[5] = awdb_functions.parse_unsigned
-- parse signed int
data_parse_type_functions[6] = awdb_functions.parse_signed
-- parse float
data_parse_type_functions[7] = awdb_functions.parse_float
-- parse double
data_parse_type_functions[8] = awdb_functions.parse_double

-- read metadata
function awdb_functions:read_meta(meta_length)
    local start_len = 2 + meta_length
    -- lua中的索引是从1开始的，所以这里的的起始位置要从3开始
    local meta_str = string.sub(self.conts, 3, start_len)

    local meta = cjson.decode(meta_str)

    local node_count = meta["node_count"]
    local byte_len = meta["byte_len"]
    local base_offset = node_count * byte_len * 2 + start_len

    meta["start_len"] = start_len
    meta["base_offset"] = base_offset

    return meta
end

-- parse awdb data (with data structures)
function awdb_functions:decode_content_structure(offset)
    local content = awdb_functions:parse_data(self, offset)
    local result = awdb_functions:map_key_value(self, content)
    return result
end

-- parse awdb data directly
function awdb_functions:decode_content_direct(offset)
    local columns = self.meta["columns"]
    local data_str = string.sub(self.conts, offset + 1, offset + 4)
    local data_len = sunpack(">I", data_str)
    offset = offset + 4
    local str = string.sub(self.conts, offset + 1, offset + data_len)

    local str_json = '["' .. string.gsub(str, "\t", '","') .. '"]'
    local values = cjson.decode(str_json)

    local result = {}
    for i = 1, #columns do
        result[columns[i]] = values[i]
    end

    return result
end

-- Search binary tree to find the location information of the corresponding IP
function awdb_functions:find_tree_index(self, bits)
    local node_index = 0
    local byte_len = tonumber(self.meta["byte_len"])
    local start_len = tonumber(self.meta["start_len"])
    local node_count = tonumber(self.meta["node_count"])

    local bit_length = #bits * 8
    local pl = 0

    while pl < bit_length and node_index < node_count do
        local b = bit.band(bits[math.floor(pl / 8) + 1], 0xFF)
        local bit_value = bit.band(1, bit.rshift(b, 7 - (pl % 8)))

        local offset = node_index * byte_len * 2 + bit_value * byte_len + start_len
        local o1, o2, o3, o4 = self.conts:byte(offset + 1, offset + byte_len)
        node_index = o1 * 16777216 + o2 * 65536 + o3 * 256 + o4

        pl = pl + 1
    end

    if node_index == node_count then
        print("Invalid node_index in search tree")
        return ""
    elseif node_index > node_count then
        return node_index
    end

    error("Invalid node_index in search tree")
end

-- parse data content
function awdb_functions:parse_content(self, offset)
    local data_type = tonumber(self.conts:byte(offset + 1))
    offset = offset + 1
    -- Find the corresponding parser function based on the data type
    local func = data_parse_type_functions[data_type]
    if func == nil then
        error("Unexpected data type number (" .. data_type .. "), please check! ")
    end

    return func(self, offset)
end

-- parse data
function awdb_functions:parse_data(self, offset)
    local content, _ = awdb_functions:parse_content(self, offset)
    return content
end

-- column value mapping
function awdb_functions:map_key_value(self, content)
    local columns = self.meta["columns"]
    if #columns == #content then
        local result = {}
        for i = 1, #columns do
            result[columns[i]] = content[i]
        end
        return result
    else
        local result = {}
        local count = #content - 1
        for i = 1, count do
            result[columns[i]] = content[i]
        end

        local multi_areas_name = columns[#columns - 1]
        local keys = columns[#columns]
        local values = content[#content]
        local tail = {}
        for _, v in ipairs(values) do
            local multi = {}
            for i = 1, #keys do
                multi[keys[i]] = v[i]
            end
            table.insert(tail, multi)
        end
        result[multi_areas_name] = tail
        return result
    end
end

-- two byte to int16
local function read_i16(b1, b2)
    assert(0 <= b1 and b1 <= 0xff) -- 8 bit
    assert(0 <= b2 and b2 <= 0xff) -- 8 bit
    local num = bit.bor(bit.lshift(b1, 8), bit.lshift(b2, 0))
    return num
end

-- read metadata length
local function meta_length(conts)

    local one, two = conts:byte(0, 2)

    return read_i16(one, two)
end

-- read metadata, return metatable
local function new_awdb(conts)
    local metaLen = meta_length(conts)

    local self = setmetatable({
        conts = conts;
        meta = nil;
    }, awdb_meta)

    self.meta = self:read_meta(metaLen)

    return self
end

-- read awdb file
local function read_awdb(curr_file)
    local f1 = io.open(curr_file, "rb")
    assert(f1, curr_file .. " File failed to open")

    local conts = f1:read("*a")

    f1:close()
    assert(conts, curr_file .. " File read failed")

    return new_awdb(conts)
end

-- converts an unsigned number to a signed number
local function unsigned_to_singed(num)
    -- judgement whether the highest bit is 1
    if num >= 128 then
        -- If the highest bit is 1, it is converted to a negative number
        num = num - 256
    end
    return num
end

-- ipv4 to array
local function ipv4_to_array(ip)
    local o1, o2, o3, o4 = ip:match("(%d%d?%d?)%.(%d%d?%d?)%.(%d%d?%d?)%.(%d%d?%d?)")
    assert(o1, "invalid IPv4 address")
    o1 = tonumber(o1, 10)
    o2 = tonumber(o2, 10)
    o3 = tonumber(o3, 10)
    o4 = tonumber(o4, 10)
    assert(o1 <= 255 and o2 <= 255 and o3 <= 255 and o4 <= 255, "invalid IPv4 address")
    return {
        unsigned_to_singed(o1);
        unsigned_to_singed(o2);
        unsigned_to_singed(o3);
        unsigned_to_singed(o4);
    }
end

-- ipv6 compressed format to ipv6 full format
local function full_ipv6(str)
    local ip_chunk = {}
    for chunk in str:gmatch("(%x%x?%x?%x?):?") do
        for i = 1, 4 - #chunk, 1 do
            chunk = '0' .. chunk
        end
        table.insert(ip_chunk, chunk)
    end
    return ip_chunk
end

-- ipv4 maps ipv6
local function ipv4_map_ipv6(ipv4)
    local ipv6_prefix = "0000:0000:0000:0000:0000:ffff:"
    local o1, o2, o3, o4 = ipv4:match("(%d%d?%d?)%.(%d%d?%d?)%.(%d%d?%d?)%.(%d%d?%d?)")
    local ipv6_suffix = string.format("%02x%02x:%02x%02x", tonumber(o1), tonumber(o2), tonumber(o3), tonumber(o4))
    return ipv6_prefix .. ipv6_suffix
end

-- ipv6 split
local function split_ipv6(str)

    if str:match('::') then
        -- ipv4 mapping ipv6 format
        if str:match('%.') then
            str = ipv4_map_ipv6(str)
        else
            -- ipv6 compressed format
            local middle = '0000'
            local ip_chunk = {}
            local ip_former, ip_after = str:match("^([%x:]-)::([%x:]*)$")
            former = full_ipv6(ip_former)
            after = full_ipv6(ip_after)
            for _, v in ipairs(former) do
                table.insert(ip_chunk, v)
            end
            
            for i = 1, 8 - #former - #after do
                table.insert(ip_chunk, middle)
            end
            
            for _, v in ipairs(after) do
                table.insert(ip_chunk, v)
            end
            str = table.concat(ip_chunk, ':')
        end
    elseif not str:match('::') and not str:match('%.') and str:len() < 39 then
        -- another IPv6 compression format
        local ip_chunk = full_ipv6(str)
        str = table.concat(ip_chunk, ':')
    end

    local components = {}
    local n = 0
    for u16 in str:gmatch("(%x%x?%x?%x?):?") do

        n = n + 1
        assert(u16, "invalid IPv6 address")
        components[n] = u16
    end

    return components
end

--ipv6 to array
local function ipv6_to_array(ip)
    local components = split_ipv6(ip)
    assert(#components == 8, "invalid IPv6 address")

    local bits = {}
    local j = 1
    for i = 1, 8 do
        local u16 = components[i]
        bits[j] = unsigned_to_singed(tonumber(u16:sub(1, 2), 16)) -- convert to a signed number
        bits[j + 1] = unsigned_to_singed(tonumber(u16:sub(3, 4), 16)) -- convert to a signed number
        j = j + 2
    end
    return bits
end

-- int to ipv4
local function int_to_ipv4(n)
    if n then
        n = tonumber(n)
        local n1 = math.floor(n / (2 ^ 24))
        local n2 = math.floor((n - n1 * (2 ^ 24)) / (2 ^ 16))
        local n3 = math.floor((n - n1 * (2 ^ 24) - n2 * (2 ^ 16)) / (2 ^ 8))
        local n4 = math.floor((n - n1 * (2 ^ 24) - n2 * (2 ^ 16) - n3 * (2 ^ 8)))
        return n1 .. "." .. n2 .. "." .. n3 .. "." .. n4
    end
    return "0.0.0.0"
end

-- verify the IP type
local function judge_ip_type(ip)
    local ip_array = { error = 0, ipv4 = 1, ipv6 = 2 }

    -- check for format 1.11.111.111 for ipv4
    local ipv4_chunks = { ip:match("^(%d+)%.(%d+)%.(%d+)%.(%d+)$") }
    if #ipv4_chunks == 4 then
        for _, v in pairs(ipv4_chunks) do
            if tonumber(v) > 255 then
                return ip_array.error
            end
        end
        return ip_array.ipv4
    end

    -- check for ipv6 format, should be 8 'chunks' of numbers/letters
    -- without leading/trailing chars
    -- or fewer than 8 chunks, but with only one `::` group
    local ipv6_chunks = { ip:match("^" .. (("([a-fA-F0-9]*):"):rep(8):gsub(":$", "$"))) }
    if #ipv6_chunks == 8
            or #ipv6_chunks < 8 and ip:match('::') and not ip:gsub("::", "", 1):match('::') then
        for _, v in pairs(ipv6_chunks) do
            if #v > 0 and tonumber(v, 16) > 65535 then
                return ip_array.error
            end
        end
        return ip_array.ipv6
    end
end

-- decode content structure
parse_type_func[1] = awdb_functions.decode_content_structure
-- decode content direct
parse_type_func[2] = awdb_functions.decode_content_direct

-- find IPv4 location information
function awdb_functions:find_ipv4_loc(ip)
    return self:find_ip_loc(ip)
end

-- find IPv6 location information
function awdb_functions:find_ipv6_loc(ip)
    return self:find_ip_loc(ip)
end

-- find IPv4 or IPv6 location information
function awdb_functions:find_ip_loc(ip)

    local ip_version = self.meta["ip_version"]
    if ip_version == "4_6" then
        if type(tonumber(ip)) == "number" then
            if tonumber(ip) < 2 ^ 32 then
                ip = int_to_ipv4(ip)
                ip = ipv4_map_ipv6(ip)
            elseif tonumber(ip) >= 2 ^ 32 then
                error("This release does not support long type queries for IPv6! The IP is : " .. ip)
            end
            -- handle this situation 192.168.1.1
        elseif judge_ip_type(ip) == 1 then
            ip = ipv4_map_ipv6(ip)
        end
    end

    if ip_version == "4" then
        if type(tonumber(ip)) == "number" then
            if tonumber(ip) >= 2 ^ 32 then
                error("The database is IPv4 library, but you are using ipv6 queries!, The IP is : " .. ip)
            end

            ip = int_to_ipv4(ip)
        end
    elseif ip_version == "6" then
        if type(tonumber(ip)) == "number" then
            if tonumber(ip) < 2 ^ 32 then
                error("The database is IPv6 library, but you are using ipv4 queries! The IP is : " .. ip)
            end

            error("This release does not support long type queries for IPv6! The IP is : " .. ip)
        end
    end

    local ip_bit
    local ip_type = judge_ip_type(ip)
    if ip_version == "4" then
        if ip_type == 1 then
            ip_bit = ipv4_to_array(ip)
        else
            error("The database is IPv4 library, but you are using ipv6 queries!, The IP is : " .. ip)
        end
    elseif ip_version == "6" or ip_version == "4_6" then
        if ip_type == 2 then
            ip_bit = ipv6_to_array(ip)
        else
            error("The database is IPv6 library, but you are using ipv4 queries! The IP is: " .. ip)
        end
    end

    local node_index = awdb_functions:find_tree_index(self, ip_bit)
    if node_index == "" then
        return {}
    end

    local node_count = tonumber(self.meta["node_count"])
    local base_offset = tonumber(self.meta["base_offset"])
    local decodeType = tonumber(self.meta["decode_type"])

    local offset = base_offset + node_index - node_count - 10

    local func = parse_type_func[decodeType]

    if func == nil then
        error("Unexpected decode type number (" .. decodeType .. "), please check! ")
    end

    return func(self, offset)
end

return {
    new_awdb = new_awdb;
    read_awdb = read_awdb;
}
